import os
import mindspore
import pickle


class EmbedingGeneratorDataset:
    def __init__(self, embedding_dir, dataset, pretraining_scheme, train, verbose):
        if train:
            self.metasplit = ['train', 'val']
        else:
            self.metasplit = ['test']
        
        self.verbose = verbose
        self.embedding_path = os.path.join(embedding_dir, dataset, pretraining_scheme)

        self.embeddings = {}
        for d in self.metasplit:
            if self.verbose:
                print('\nLoading data from ' + os.path.join(self.embedding_path, d+'_embeddings.pkl') + '...')
            self.embeddings[d] = pickle.load(open(os.path.join(self.embedding_path, d+'_embeddings.pkl'), 'rb'), encoding='latin1')
        
        # sort images by class
        self.image_by_class = {}
        self.embed_by_name = {}
        self.class_list = {}
        for d in self.metasplit:
            """依据类对图片进行分类"""
            self.image_by_class[d] = {}
            self.embed_by_name[d] = {}
            self.class_list[d] = set()
            keys = self.embeddings[d]["keys"]
            """获得文件的下标i和文件名k"""
            for i, k in enumerate(keys):
                """分割得到类名，图片名"""
                _, class_name, img_name = k.split('-')
                """出现未收录的类则创建一个新类"""
                if (class_name not in self.image_by_class[d]):
                    self.image_by_class[d][class_name] = []
                """依据类对图片进行分类,共20类"""
                self.image_by_class[d][class_name].append(img_name)
                """匹配每个图片的特征向量"""
                self.embed_by_name[d][img_name] = self.embeddings[d]["embeddings"][i]
                # construct class list
                """生成类的字典"""
                self.class_list[d].add(class_name)

            """将字典转化为列表"""
            self.class_list[d] = list(self.class_list[d])
            if self.verbose:
                print('\nFinish constructing ' + d + ' data, total %d classes.' % len(self.class_list[d]))

    def __len__():
        pass

    def __getitem__():
        pass